import sys
sys.path.append("/home/lixiaoyu/vscode_files/nano-vllm")
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# os.environ["CUDA_VISIBLE_DEVICES"] = "1"  #代表只使用第3个gpu
from datetime import datetime
from nanovllm import LLM, SamplingParams
from nanovllm.engine.spec_decode.utils import write_list_to_json, read_list_from_json
from transformers import AutoTokenizer
from omegaconf import ListConfig, DictConfig
from typing import List, Union, Optional
import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader, SequentialSampler
from torchdata.stateful_dataloader import StatefulDataLoader
from transformers import AutoTokenizer, PreTrainedTokenizer, ProcessorMixin

import torch_functional as verl_F

current_time = datetime.now()
context_info = f"当前时间是{current_time}\n"
sys_prompt = context_info + """
以下内容禁止输出给用户

你是一个差旅助手，叫作小滴，由L-lab开发。对于和差旅无关的问题，要礼貌的拒绝回答

调用工具存在依赖关系，所以先调用一个工具，根据工具返回结果，进一步思考，调用下一个工具，一步接一步。
遵循以下形式：
...
<tool_call>...</tool_call>
<tool_response>...</tool_response>
<tool_response_thinking>...</tool_response_thinking>
...
<tool_call>...</tool_call>
<tool_response>...</tool_response>
<tool_response_thinking>...</tool_response_thinking>


航班查询工具是一个组合工具，该工具整合了poi查询和打车查询的功能，当用户需要大小交通衔接时可直接设置depart_poi或者arrival_poi，使用方法如下：
flight_search(departure_city_name,arrival_city_name,depart_date,depart_poi, arrival_poi,**kwargs)
departure_city_name:必填参数，出发地所在城市名，如“上海”
arrival_city_name:必填参数，到达地所在城市名(只支持国内城市，不支持境外城市名)，如“杭州”
depart_date:出发日期，如"2025-01-03"
depart_poi:必须属于departure_city_name城市，用户从该poi（poi名称即可，不需要详细地址）出发前往机场乘机，必填参数，用户没有明确表达时，设置成departure_city_name
arrival_poi:必须属于arrival_city_name城市，用户下飞机后想要前往的poi地点（poi名称即可，不需要详细地址），必填参数，用户没有明确表达时，设置成arrival_city_name
is_transfer:是否中转，如果要查询中转航班，设置成1，搜索直飞，设置成0，必填参数
depart_time_earliest:用户最早从depart_poi出发时间（需要带日期），如"2025-01-01 07:00"
arrive_time_latest:用户最晚的到达arrival_poi的时间（需要带日期），如"2025-01-02 06:00"，当用户对最晚到达时间有要求时，使用该参数筛选航班
transport_depart_time_earliest:大交通(飞机/火车)的最早出发时间（需要带日期）
transport_depart_time_latest:大交通的最晚出发时间（需要带日期）
transport_arrive_time_latest:大交通的最晚到达时间（需要带日期）
airline_name:可选参数，填入指定的航司名称（必须规范使用4个字的简称），支持多家航司，如"中国国航,东方航空"
departure_airport_name:起飞机场，可选参数，用户对出发机场有要求设置该参数
arrival_airport_name:降落机场，使用四字简称，如“首都机场”，“虹桥机场”，可选参数，用户对到达机场有要求设置该参数
cabin_types:舱位，支持头等舱/商务舱/经济舱，可选参数
min_sale_price:最低售价，单位：元，可选参数
max_sale_price:最高售价，单位：元，可选参数
has_pre_trip_transport_on_the_same_day:当天包含前置的大交通时，需要把该参数设置成1；同时需要把depart_time_earliest设置成上一个交通方式的到达时间之后。
使用时遵循以下格式：<tool_call>flight_serach(**kwargs)</tool_call>。kwargs只可以传入上面列举的参数，不可用其他。
举例：2025年7月2号8点要从北京大学出发去上海，坐飞机。
可以这么调用工具查询满足时间要求的机票:<tool_call>flight_search(departure_city_name="北京", arrival_city_name="上海", depart_date="2025-07-02", depart_poi="北京大学", arrival_poi="上海", depart_time_earliest="2025-07-02 08:00")</tool_call>

火车票查询工具是一个组合工具，可以查询两个站点间的直达和中转信息；该工具整合了poi查询和打车查询的功能，当用户需要大小交通衔接时可直接设置depart_poi或者arrial_poi，使用方法如下：
train_search(departure_city_name, arrival_city_name, depart_date, depart_poi, arrival_poi, **kwargs)
departure_city_name:出发城市，如"北京"，必填参数
arrival_city_name:到达城市，如"上海"，必填参数
depart_date:出发日期，如"2025-01-01"，必填参数
depart_poi:必须属于departure_city_name城市，用户从该poi（poi名称即可，不需要详细地址）出发前往车站乘车，必填参数，用户没有明确表达时，设置成departure_city_name
arrival_poi:必须属于arrival_city_name城市，用户下火车后想要前往的poi地点（poi名称即可，不需要详细地址），必填参数，用户没有明确表达时，设置成arrival_city_name
depart_station:出发站名，如"北京","北京西"，可选参数，用户对出发站有要求设置该参数
arrive_station:到达站名，如"武汉"，可选参数，用户对到达站有要求设置该参数
is_transfer:是否中转，如果要查询中转列车，设置成1，搜索直达，设置成0，必填参数
train_type:火车类型，支持：高铁/动车/普通车，多个选择用","连接，注意区分高铁和动车，用户没提到不需填写
depart_time_earliest:用户最早从depart_poi出发时间（需要带日期），如"2025-01-01 07:00"，不需要累加打车的时间
arrive_time_latest:用户最晚的到达arrival_poi的时间（需要带日期），如"2025-01-02 06:00"，当用户对最晚到达时间有要求时，使用该参数筛选火车
transport_depart_time_earliest:大交通(飞机/火车)的最早出发时间（需要带日期）
transport_depart_time_latest:大交通的最晚出发时间（需要带日期）
transport_arrive_time_latest:大交通的最晚到达时间（需要带日期）
max_price:票价不超过的最大价格，单位元
seat_name:座席类型，支持的坐席有这些：{商务座,一等座,二等座,硬座,软座,硬卧,软卧}，用户没提到不需填写
has_pre_trip_transport_on_the_same_day:当天包含前置的大交通时，需要把该参数设置成1；同时需要把depart_time_earliest设置成上一个交通方式的到达时间之后。
middle_station_name:中转的车站，期望从哪里中转，is_transfer=1时候才可以使用。
对于中转查询，使用is_transfer参数，无需自己分段查询。
举例：
直达查询：<tool_call>train_search(departure_city_name="北京", arrival_city_name="南京", depart_date="2025-05-01",depart_poi="北京大学", arrival_poi="南京人民医院",depart_time_latest="2025-05-01 14:00")</tool_call>
中转查询：<tool_call>train_search(departure_city_name="北京", arrival_city_name="南京", depart_date="2025-05-01",depart_poi="北京大学", arrival_poi="南京人民医院",depart_time_latest="2025-05-01 18:00",is_transfer=1)</tool_call>
举例：2025年7月10号从先北京去广州，再从广州去中山。
其中广州到中山的交通查询应该这么调用工具:<tool_call>train_search(departure_city_name="广州", arrival_city_name="中山", depart_date="2025-07-10",depart_poi="广州", arrival_poi="中山", has_pre_trip_transport_on_the_same_day=1, depart_time_earliest=$上一段航班的结束时刻+$预留换乘buffer)</tool_call>

机火同查工具整合了航班查询工具和火车票查询工具，可以同时查询机票和火车票；当用户没有对交通方式有要求的时候，强烈推荐使用该工具
flight_and_train_search(departure_city_name, arrival_city_name, depart_date, depart_poi, arrival_poi, **kwargs)
支持所有航班查询工具和火车票查询工具的参数
使用时遵循以下格式：<tool_call>flight_and_train_search(**kwargs)</tool_call>。
举例：2025年7月1号，12点要从武汉大学出发去上海滴滴百川大厦开会，会议时间是晚上8点，要及时赶上。
可以这么调用工具查询满足时间要求的机票和火车票：<tool_call>flight_and_train_search(departure_city_name="武汉", arrival_city_name="上海", depart_date="2025-07-01", depart_poi="武汉大学", arrival_poi="滴滴百川大厦", depart_time_earliest="2025-07-01 12:00", arrive_time_latest="2025-07-01 20:00")</tool_call>

酒店查询工具是一个函数调用，使用方法如下：
hotel_search(city_name,checkin_date,checkout_date,**kwargs)
city_name:城市名称(只支持国内城市，不支持境外城市名)，如“北京”，必填参数
checkin_date:入住日期，如"2025-01-01",必填参数，用户未指明的话，默认用当前时间的日期，必填参数
checkout_date:离店日期（至少比入店日期晚一天），如"2025-01-03",必填参数，用户未指明的话，默认用当前时间的日期加一天，必填参数
hotel_name:酒店名称，用户没提到不需填写，只支持单个酒店查询，如果用户提到多个指定酒店，需要依次调用工具查询，不要一次性查询多个酒店
brand_name:酒店品牌名，如“汉庭”，“如家”，可选参数，用户没提到不需填写，可以支持多个品牌查询，用逗号分隔，例如"汉庭,如家"
room_name:房型，支持：大床房/双床房/家庭房/套房，用户没提到不需填写
landmark:地标，如“故宫”，如果用户提到xx附近的酒店，直接把landmark设置成xx，用户如果没提到则不需要填写
distance_range:酒店到landmark的距离，单位：米，用户如果没提到不需要填写，如果提到是附近，请设置成5000
price_low:预算下限，可选参数，用户没提到不需填写
price_high:预算上限，可选参数，用户没提到不需填写
level_name:星级，用户没提到不需填写
has_breakfast:是否包含早餐，1：包含，0:不包含，用户没提到不需填写
has_window:房间是否有窗户，1:有窗，用户没提到不需填写
conforming_quota:设置成1则筛选差标内的；可选参数，用户明确提到需满足差标时候才设置，否则不传
使用时遵循以下格式：<tool_call>hotel_search(**kwargs)</tool_call>。kwargs只可以传入上面列举的参数，不可用其他。
如：<tool_call>hotel_search(city_name="杭州", checkin_date="2025-05-01", checkout_date="2025-05-03", landmark="浙江大学")</tool_call>

生成卡片(只对航班/火车/酒店)，卡片信息（除了is_last字段，is_last用来标记同种类型推荐列表的最后一个）必须保证来自于tool_response，尤其是item_key字段必须保证和tool_response中对应资源的字段完全一致。
如果涉及到中转的火车查询，根据工具结果生成一张卡片即可。
对于火车，<card>{"category":"火车","train_code":,"from_station_name":,"to_station_name":,"depart_date":,"item_key":,"is_last":1or0(是否是推荐的最后一个火车)}</card>
对于航班，<card>{"category":"航班","departure_datetime":,"flight_number":,"departure_airport_name":,"arrival_airport_name":,"item_key":,"is_last":1or0(是否是推荐的最后一个航班)}</card>
对于酒店，<card>{"category":"酒店","hotel_name":,"checkin_date":,"checkout_date":,"item_key":,"is_last":1or0(是否是推荐的最后一个酒店)}</card>
注意：<card>...</card>也只能在答案中使用（在</think>后面），禁止出现在思考过程中

你具有联网搜索的能力。在处理一些实时信息时候，你可以通过使用联网搜索功能来获取相关信息，然后给用户更精准的回复。
联网搜索是一个函数调用：
web_search(query)
query:搜索的关键词，要做到言简意赅。
使用时遵循以下格式<tool_call>web_search(**kwargs)</tool_call>,如<tool_call>web_search(query="2025年5月1日 北京 天气")</tool_call>
注意机票/酒店/火车的相关查询，使用专门的机票/酒店/火车查询工具，不要使用web_search。
注意不要重复搜索已经搜索过的关键词。

你具有poi搜索的能力，可以根据地点描述查询地址。
poi搜索是一个函数调用：
poi_search(query,city_name,**kwargs)
query:地点的描述
city_name:地点所在的城市名
is_area:如果地点属于一个范围，设置is_area=1,否则无需设置，可选参数
使用时遵循以下格式<tool_call>poi_search(**kwargs)</tool_call>,如<tool_call>poi_search(query="北邮",city_name="北京")</tool_call>
注意poi_search结果可能会有多个，如果不确定是哪一个，可以把列表告诉用户，让用户从列表中选择。

你具有滴滴打车查询的能力，可以获得两个地点之间的距离，打车时间的信息。
打车查询是一个函数调用：
ride_hailing(origin, destination, city_name)
origin:出发的地点
destination:到达的地点
city_name:查询的城市名
使用时遵循以下格式<tool_call>ride_hailing(**kwargs)</tool_call>，如<tool_call>ride_hailing(origin="故宫", destination="颐和园", city_name="北京")</tool_call>
注意该工具只能用于同一个城市的两个地点之间的路线查询。

在所有其他工具调用结束后必须调用整合资源工具，用于后续生成卡片。使用方法如下：
整合资源()
使用时遵循以下格式：<tool_call>整合资源()</tool_call>
注意：该工具无需传入参数

你会生成卡片，卡片生成的方法：
对于火车，<card>{"category":"火车","depart_date":,"TrainName":,"StartStationName":,"EndStationName":,"StartTime":,"EndTime":}</card>
对于航班，<card>{"category":"航班","departure_datetime":,"flight_number":,"departure_airport_name":,"arrival_airport_name":,"cabin_name":}</card>
对于酒店，<card>{"category":"酒店","didi_hotel_id":,"hotel_name":,"checkin_date":,"checkout_date":}</card>
举例说明如下：
推荐入住xx酒店，酒店信息介绍，
<card>...</card>
推荐xx航班，航班信息介绍，
<card>...</card>
..

日期时间助手是一个函数调用，对于复杂的日期时间推理可以使用它
使用方法
copilot(question)
question是问题的描述,要描述的清晰简洁
使用时遵循以下格式<tool_call>copilot(**kwargs)</tool_call>
如：<tool_call>copilot(question="现在是2025-04-01,下周三的日期是多少")</tool_call>

需要遵循的原则：
对于不确定信息，不要自己假设，要主动询问用户。
推荐酒店/火车/航班采用markdown格式，日期和其他信息要清晰，并且要生成相应的卡片。
推荐的酒店/火车/航班必须来自tool_response中的结果，禁止假设，不要遗漏工具的调用。
使用工具时一定要以<tool_call>为开头。
如果用户涉及到准时的事件，行程规划必须保证预留足够的时间，预留多久时间需要推理，符合常识。
默认至少要给出去程和返程的交通方式（航班或者火车），酒店住宿（酒店推荐两个候选）的方案。
最终规划的行程结果放到<answer>后输出。
"""

model_path = "/home/lixiaoyu/huggingface/Qwen3-0.6B/"
# model_path = "/home/lixiaoyu/huggingface/Qwen3-8B/"

prompts = [
    "我有从北京到广州的行程，帮我安排出行，我希望在明天晚上出发。",
    "我有从广州到石家庄的行程，帮我安排出行，我希望在7/16晚上出发。",
]


def compute_position_id_with_mask(mask):
    return torch.clip(torch.cumsum(mask, dim=-1) - 1, min=0, max=None)

def collate_fn(data_list: list[dict]) -> dict:
    tensors = {}
    non_tensors = {}

    for data in data_list:
        for key, val in data.items():
            if isinstance(val, torch.Tensor):
                if key not in tensors:
                    tensors[key] = []
                tensors[key].append(val)
            else:
                if key not in non_tensors:
                    non_tensors[key] = []
                non_tensors[key].append(val)

    for key, val in tensors.items():
        tensors[key] = torch.stack(val, dim=0)

    for key, val in non_tensors.items():
        non_tensors[key] = np.array(val, dtype=object)

    output = {}
    output.update(tensors)
    output.update(non_tensors)
    return output
class RLHFDataset(Dataset):
    """
    We assume the dataset contains a column that contains prompts and other information
    """

    def __init__(self,
                 data_files: Union[str, List[str]],
                 tokenizer: PreTrainedTokenizer,
                 prompt_key='prompt',
                 max_prompt_length=1024,
                 filter_prompts=True,
                 cache_dir='~/.cache/verl/rlhf',
                 chat_template_func=None,
                 return_raw_chat=False,
                 truncation='left',
                 processor: Optional[ProcessorMixin] = None,
                 config: DictConfig=any):
        parquet_files = data_files
        if not isinstance(parquet_files, (List, ListConfig)):
            parquet_files = [parquet_files]

        self.parquet_files = parquet_files
        self.cache_dir = os.path.expanduser(cache_dir)
        self.tokenizer = tokenizer
        self.processor = processor
        self.config = config

        self.prompt_key = prompt_key
        self.max_prompt_length = max_prompt_length
        self.filter_prompts = filter_prompts

        self.return_raw_chat = return_raw_chat
        self.chat_template_func = chat_template_func
        self.truncation = truncation

        self._download()
        self._read_files_and_tokenize()

    def _download(self):
        from fs import copy_local_path_from_hdfs
        for i, parquet_file in enumerate(self.parquet_files):
            self.parquet_files[i] = copy_local_path_from_hdfs(src=parquet_file, cache_dir=self.cache_dir)

    def _read_files_and_tokenize(self):
        dataframes = []
        for parquet_file in self.parquet_files:
            # read parquet files and cache
            dataframe = pd.read_parquet(parquet_file)
            dataframes.append(dataframe)
        self.dataframe = pd.concat(dataframes)

        print(f'original dataset len: {len(self.dataframe)}')

        # filter out too long prompts
        tokenizer = self.tokenizer
        prompt_key = self.prompt_key
        # self.dataframe = self.dataframe[self.dataframe.apply(lambda doc: len(
        #     tokenizer.apply_chat_template(doc[prompt_key], add_generation_prompt=True)) <= self.max_prompt_length,
        #                                                      axis=1)]

        print(f'filter dataset len: {len(self.dataframe)}')

    def __len__(self):
        return len(self.dataframe)

    def __getitem__(self, item):
        """
        Note that we also return the raw_input_ids so that it can be combined with other chat template
        """
        row_dict = self.dataframe.iloc[item].to_dict()
        chat = row_dict.pop(self.prompt_key)
        raw_prompt_key = "raw_prompt"
        raw_prompt = row_dict.pop(raw_prompt_key)
        # print(f'raw: {raw_prompt}')
        line = [{"role":"system", "content": sys_prompt}, {"role": "user", "content": raw_prompt}]
        prompts = [self.tokenizer.apply_chat_template(line, tokenize=False, add_generation_prompt=True)]

        # prompt_with_chat_template = chat
        input_ids, attention_mask = verl_F.tokenize_and_postprocess_data(prompt=prompts[0],
                                                                         tokenizer=self.tokenizer,
                                                                         max_length=self.max_prompt_length,
                                                                         pad_token_id=self.tokenizer.pad_token_id,
                                                                         left_pad=True,
                                                                         truncation=self.truncation)

        position_ids = compute_position_id_with_mask(attention_mask)

        row_dict['input_ids'] = input_ids[0]
        row_dict['attention_mask'] = attention_mask[0]
        row_dict['position_ids'] = position_ids[0]

        # encode prompts without chat template
        raw_input_ids, _ = verl_F.tokenize_and_postprocess_data(prompt=raw_prompt,
                                                                         tokenizer=self.tokenizer,
                                                                         max_length=self.max_prompt_length,
                                                                         pad_token_id=self.tokenizer.pad_token_id,
                                                                         left_pad=True,
                                                                         truncation=self.truncation)
        
        row_dict['raw_prompt'] = raw_input_ids[0]
        row_dict['sys_prompt'] = sys_prompt

        # add index for each prompt
        row_dict["extra_info"] = {"index": item }
        return row_dict

def create_dataloader(train_files, tokenizer, batch_size=100):
    dataset = RLHFDataset(data_files=train_files,tokenizer=tokenizer, max_prompt_length=4096)
    sampler = SequentialSampler(data_source=dataset)

    dataloader = StatefulDataLoader(dataset=dataset,
                                    batch_size=batch_size,
                                    num_workers=4,
                                    drop_last=True,
                                    collate_fn=collate_fn,
                                    sampler=sampler)

    assert len(dataloader) >= 1
    print(f'Size of train dataloader: {len(dataloader)}')

    return dataloader

def _pre_process_inputs(pad_token_id, prompt_token_ids: torch.Tensor) -> List[int]:
    # remove the left padding in the prompt token_id
    # pad_token_id = self.llm_engine.tokenizer.pad_token_id if self.llm_engine.tokenizer.pad_token_id is not None else self.llm_engine.tokenizer.eos_token_id
    non_pad_index = torch.nonzero(prompt_token_ids != pad_token_id, as_tuple=False)[0][0]
    token_ids = prompt_token_ids[non_pad_index:].tolist()
    return token_ids
def preprocess(prompts, tokenizer):
    idx = prompts['input_ids']  # (bs, prompt_length)

    batch_size = idx.size(0)

    idx_list = []
    # parse idx from torch.Tensor to List[List[str]]
    for i in range(batch_size):
        idx_list.append(_pre_process_inputs(tokenizer.pad_token_id, idx[i]))
    return idx_list

if __name__ == "__main__":
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    dataloader = create_dataloader("/home/lixiaoyu/vscode_files/nano-vllm/speculative_decoding/train_0630_simple.parquet", tokenizer, 32)    
    sampling_params = SamplingParams(temperature=0.8, max_tokens=1024, ignore_eos=False)

    ngram_datasets = read_list_from_json("/home/lixiaoyu/vscode_files/nano-vllm/speculative_decoding/logs_8b.json")
    ngram_datasets = [tokenizer.encode(seq) for seq in ngram_datasets]

    llm = LLM(
        model=model_path,
        enforce_eager=True, 
        tensor_parallel_size=1,
        gpu_memory_utilization=0.8,
        max_model_len=8092,
        speculative_config={"method": "ngram", "num_speculative_tokens": 7},
        ngram_datasets=ngram_datasets
    )
    
    outputs_list = []
    for i, data in enumerate(dataloader):
        input_ids = preprocess(data, tokenizer)
        prompts = [tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids]
        outputs = llm.generate(prompts, sampling_params)
        outputs_list += [output['text'] for output in outputs]
        
        if i == 4:
            break
    
    write_list_to_json(outputs_list, "./logs.json")
    
    # print(f"Completion: {outputs[0]['text']}")
    # for prompt, output in zip(prompts, outputs):
    #     print("\n")
    #     # print(f"Prompt: {prompt!r}")
    #     print(f"{output['text']!r}")
    
        
    